/**
  SUNNY project, Anyang Normal University, IMP-CAS
  \class TALanczos
  \brief Lanczos algorithm for solving large sparse matrix iteratively.
  Transcribbed from momd/TADiagonalize::Lanczos and momd/TADiagonalize::LanczosPurge,
  but with new classes representing large sparse matrix using ROOT - TTree object
  storing large sparse matrix and vectors, so as to save memory. This is a static
  class.
  Dedicated for diagonalization problems arising from large-scale shell model
  calculations.
  \author SUN Yazhou, asia.rabbit@163.com
  \since 2021/02/12, Chinese Lunar New Year
  \date Last modified: 2021/02/12 by SUN Yazhou
  \copyright 2020-2021 SUN Yazhou
  \copyright SUNNY project, Anyang Normal University, IMP-CAS
*/

#include <vector>
#include <algorithm>
#include "config.h" // contains user switch of definition of ROOT and OpenCV
#ifdef CATCH_CNOK
#include "catch.h"
#endif
#include "TARandom.h"
#include "TALanczos.h"
#include "TADiagonalize.h"
#include "TAException.h"
#include "TAHamiltonian.h"
#include "TASparseVec.h"
#include "TASparseMatrix.h"
#include "TAMath.h"
#include "TAHPairing.h"
#include "TASPStateManager.h"
#include "TAMBSDManager.h"
#include "TAParaManager.h"

using std::vector;
using std::min;
using std::max;

// Lanczos algorithm is devised for a few largest (in modulus) eigenvalues
// It projects H to a space with much lower dimension (Krylov subspace) as a
// tridiagonal matrix, whose eigenvalues and eigenvectors are good approximation
// of A's.
// Given input matrix a[0..n-1][0..n-1], the routine returns the first few
// (usually <=10) eigenvalues in d and the corresponding eigenvectors in z
// d should be at least of length 50, for it is drafted in the program
void TALanczos::Lanczos(TAHamiltonian &H, double *d, TASparseMatrix &z,
    TASparseVec *x, int nPair){
  static const double EPS = 1.e-6; // eigenvalue accuracy criterion
  static const int NMAX = 50; // maximum number of iterations
  if(nPair <= 0) TAException::Error("TALanczos",
    "Lanczos: illegal nPair: %d", nPair);
  // if(nPair > NMAX) TAException::Error("TALanczos",
  //   "Lanczos: nPair too large: %d", nPair);
  const int nm = min(nPair, NMAX), np = nm + 1;
  double e[np]; vector<double> alpha, beta; // e: subdiagonal, with e[0] arbitray
  // prepare for the start of the iterations //
  TASparseVec v("v"), q("q"), r("r"); // q_{j-1}, qj, q_{j+1}, 3 basis vectors
  // initialize q to (1,0,0,..)
  if(x && x->GetEntries()) q = *x; else q.Fill(0, 1.);
  q.normalize();
  TASparseMatrix Q; Q.PushBackColumn(q); // to store the Lanczos vectors
  H.DotProduct(q, r); // r = H*q = |H|qj>
  alpha.push_back(r.DotProduct(q)); // alpha = <qj|H|qj>
  r.SelfAdd(-alpha[0], q); // r -= alphaj*qj
  beta.push_back(r.norm()); // the first off-diagonal of Tk
  if(fabs(beta[0]) < EPS){ // the trial (initial) vector is an eigenvector
    d[0] = alpha[0];
    z.Clear(); z.PushBackColumn(q);
    z.Print(); // DEBUG
    return;
  } // end if
  // OK, here we commence the iteration //
  for(int j = 1; j <= nm; j++){
    v = q; // q_{j-1}
    r.Scale(1./beta[j-1],q); // normalize qj in q
    Q.PushBackColumn(q);
    H.DotProduct(q, r); // r = H*qj
    r.SelfAdd(-beta[j-1], v); // r -= beta_{j-1}*q_{j-1};
    alpha.push_back(r.DotProduct(q)); // alpha_j=<qj|H|qj>
    r.SelfAdd(-alpha[j], q); // r -= alphaj*qj;
    r.Purify(Q); // r -= Q*Q^T*r, correct roundoff error for loss of orthogonality
    beta.push_back(r.norm()); // the next beta, equals |r|
    const int jj = j+1;
    for(int i = jj; i--;) d[i] = alpha[i]; // diagonal of Tk
    for(int i = j;  i--;) e[i+1] = beta[i]; // e[0] is not used
    matrix s(jj,jj); s = 1.; // to store the eigenvector matrix of Tk
    TADiagonalize::TridiagQLImplicit(d,e,jj,s); // diagonalize Tk
    double epsilon = 0.; // check convergence
    for(int i = jj; i--;) epsilon = max(epsilon, fabs(beta[j]*s[j][i]));
    if(epsilon <= EPS || nm == j){
      TADiagonalize::EigenSort(d,s,jj);
      Q.DotProduct(s, z); // z = Q*s
      // r.Print(); r.Print(); r.Save(); return; // DEBUG
      if(nm == j) break; else return;
    } // end if
  } // end for over j
  TAException::Info("TALanczos", "Lanczos: Too many iterations occurred.");
} // end member function Lanczos

// Lanczos method featuring Ritz-pairs purging (thick restart). After getting a
// set of Ritz vectors from a plain Lanczos, we keep a smaller set and remove
// (purge) the rest. Then develop Krylov space from the new Lanczos basis. So
// unwanted directions are removed from the initial vector for each iteration.
// NOTE that the largest eigenvalues are kept.
// \param nPair: number of eigenpairs wanted.
// \param x: the initial vector to initiate the iteration.
// Still the function will return as long as encountering an invariant space,
// even before finding nPair eigenpairs. For this drawback, users are recommended
// with member method PurgeRestart(...)
// Ref.: http://people.inf.ethz.ch/arbenz/ewp/Lnotes/lsevp.pdf, p212
// Ref.: K. Wu and H. D. Simon,SIAM J. Matrix Anal. Appl., 22 (2000), pp. 602–616.
// zEx: external eigenvector set against which r is to be purified
// -- a remedy to full reorthogonalization.
void TALanczos::LanczosPurge(TAHamiltonian &H, double *d, TASparseMatrix &z,
    TASparseVec *x, int nPair){
  static const double EPS = 1.e-8; // eigenvalue accuracy standard
  static const int NMAXRST = 100; // maximum restart times
  // OFST: Nof purged Ritz vecs per restart
  static const int OFST = 3;
  // number of wanted eigenpairs and dimension for the plain Lanczos
  if(nPair <= 0) TAException::Error("TALanczos",
    "LanczosPurge: illegal nPair: %d", nPair);
  // J, K: number of the wanted eigenpairs and candidates
  const int n = int(H.GetNBasis());
  if(nPair >= n) nPair = max(2, n - OFST);
  const int J = nPair, K = J + 1 + OFST;
  // we deal with a little bit larger matrix
  if(n <= K) return Lanczos(H,d,z,x,nPair);
  const int NMAX = 20; // maximum number of iterations
  const int nm = min(n, NMAX);
  const int np = nm + 1; // maximum dimension of the final Tk
  // XXX VLA: varialble length array; FIXME //
  double e[np]; for(auto &t : e) t = 0.;
  vector<double> alpha, beta;
  int restartCnt = 0; // restart times
  bool restarted = false; // if has just restarted
  // prepare for the start of the iterations //
  TASparseVec v("v"), q("q"), r("r"); // q_{j-1}, qj, q_{j+1}, 3 basis vectors
  // initialize q to (1,0,0,..)
  if(x && x->GetEntries()) q = *x; else q.Fill(0, 1.);
  q.normalize();
  TASparseMatrix Q; Q.PushBackColumn(q); // to store the Lanczos vectors
  H.DotProduct(q, r); // r = H*q
  alpha.push_back(r.DotProduct(q)); // alpha = <qj|H|qj>
  r.SelfAdd(-alpha[0], q); // r = A*q-alphaj*qj
  beta.push_back(r.norm()); // the first off-diagonal of Tk
  // FIXME: if the trial vector is an eigenvector, return
  if(fabs(beta[0]) < EPS){
    d[0] = alpha[0];
    z.Clear(); z.PushBackColumn(q);
    return;
  } // end if
  matrix st(J+1,J+1); st = 1.; // st: the eigenvector matrix of Tk, st set to E
  // commence the iteration //
  // j: n of Lanczos vectors, so it'd be always far from nm
  for(int j = 1; j <= nm; j++){
    v = q; // q_{j-1}
    r.Scale(1./beta[j-1], q); // normalize r and store it in q
    H.DotProduct(q, r); // r = H*qj
    r.SelfAdd(-beta[j-1], v); // r -= beta_{j-1}*q_{j-1};
    alpha.push_back(r.DotProduct(q)); // alpha=<qj|H|qj>
    r.SelfAdd(-alpha[j], q); // r -= alphaj*qj;
    const int jj = j+1;
    // restoring d, e and s is nontrivial, for they are changed per iteration
    for(int i = jj; i--;) d[i] = alpha[i]; // diagonal of Tk
    for(int i = j; i--;) e[i+1] = beta[i]; // e[0] is not used
    // reconstruct Lanczos bases from the deflated Ritz pair set //
    if(restarted){ // just restarted, j = J+1 now
      // tridiagonalize the mutant (arrow) Tk to extract the Lanczos basis
      // the Householder matrix is stored in st with st.e_{J+1}=e_{J+1}
      // ensured as TridiagArrow(..) starts the reduction from the last row
      // note that e[J+1] is not used, nor is e[0]
      TADiagonalize::TridiagArrow(d,e,j,st);
      for(int i = j; i--;) alpha[i] = d[i]; // Tk recovered to tridiagonal, j=J+1
      // beta[J] (i.e. e[J+1]) doesn't belong to Tk
      for(int i = J; i--;) beta[i] = e[i+1];
      // store the Lanczos basis of the deflated Krylov space temporarily in z
      Q.DotProduct(st, z); // z = Q.st
      for(int i = J; i--;) *Q[i] = *z[i]; // accept the new Lanczos basis
      restarted = false; // restart has been dealt with
    } // end if, we can now pretend that no restart whatsoever has ever happened
    Q.PushBackColumn(q); // confirm new Lanczos basis qj
    r.Purify(Q); // r -= Q*Q^T*r, correct roundoff error for loss of orthogonality
    beta.push_back(r.norm()); // betaj=|r|
    matrix s(jj,jj); s = 1.; // stores the eigenvectors of Tk
    TADiagonalize::TridiagQLImplicit(d,e,jj,s); // diagonalize Tk
    // sort eigenpairs by eigenvalues in descending order
    TADiagonalize::EigenSort(d,s,jj);
    double epsilon = 0.; // check for convergence
    for(int i = min(jj,J); i--;) epsilon = max(epsilon, fabs(beta[j]*s[j][i]));
    // cout << "epsilon: " << epsilon << endl; // DEBUG
    if(epsilon <= EPS || nm == j){
      TADiagonalize::EigenSort(d, s, jj);
      Q.DotProduct(s, z);
      if(nm == j) break; else return;
    } // end if
    // for J 'good' Ritz pairs, we now have K candidates. K-J 'bad' ones would //
    // be identified and removed (deflated, or purged) here //
    // K Ritz pairs accumulated
    if(K-1 == j && !restarted && restartCnt < NMAXRST){
      // deflate J+1 ~ K Ritz-pairs (with the least eigenvalues) //
      // and update the tridiagonal matrix //
      alpha.erase(alpha.begin()+J, alpha.end());
      // beta_{K} reserved for future use
      beta.erase(beta.begin()+J, beta.end()-1);
      Q.DotProduct(s, z); // z=Q*s: calculate all the Ritz vectors
      Q.EraseColumn(J, Q.size()); // remove the unwanted Ritz vectors
      for(int i = J; i--;) *Q[i] = *z[i]; // change qi for zi in Q
      r.Scale(1./beta[J], q); // compute q_{K+1} = r/|r|
      Q.PushBackColumn(q); // add q_{K+1} to Qgood
      H.DotProduct(q, r); // r = H*q_{K+1}
      for(int i = 0; i < J; i++){
        alpha[i] = d[i]; // the diagonal of Tk(0:J-1, 0:J-1)
        beta[i] = beta[J]*s[j][i]; // sigma_i: <zi|H|q_{K+1}>
        r.SelfAdd(-beta[i], *z[i]);
      } // end for over i
      alpha.push_back(r.DotProduct(q)); // alpha = <q_{K+1}|H|q_{K+1}>
      r.SelfAdd(-alpha[J], q); // remove q_{K+1} component, q_{K+2} constructed
      beta[J] = r.norm(); // r is clean now, beta_{K+1}=|r_{K+1}|
      j -= K-J; // one more thing, squeeze the loop for the purge
      // j++->j=J: restarting itself moves a step forward in the Lanczos iteration
      j++;
      restarted = true; restartCnt++;
      TAException::Info("TALanczos", "LanczosPurge: Restart #%d.", restartCnt);
    } // end if(j == K)
    if(restartCnt >= NMAXRST)
      TAException::Info("TALanczos", "LanczosPurge: Too many restarts occurred.");
  } // end for over j
  TAException::Info("TALanczos", "LanczosPurge: Too many iterations occurred.");
} // end member function LanczosRestart

// restart LanczosPurge with a new initial vector orthogonal to the existing
// invariant space, with the previous eigenpairs stored.
// \param nPair: number of eigenpairs wanted
void TALanczos::PurgeRestart(TAHamiltonian &H, double *d, TASparseMatrix &z,
    int nPair){
  // temporary eigenvector container for each call of LanczosPurge
  TASparseMatrix zz;
  TASparseVec x; z.Clear();
  if(nPair <= 0) TAException::Error("TALanczos",
    "PurgeRestart: illegal nPair: %d", nPair);
  // temporary eigenvalue container for each call of LanczosPurge
  const int np = nPair + 10;
  double *dd = new double[np];
  for(int i = np; i--;) dd[i] = -9999.;

  const int n = int(H.GetNBasis());
  while(int(z.size()) < nPair){
    InitializeVector(x, z, n); // assign the initial Lanczos vector
    LanczosPurge(H,dd,zz,&x, 3);
    // stash the liberated eigenpairs
    const int nczz = zz.size(), ncz = z.size();
    for(int i = 0; i < nczz; i++){
      d[ncz + i] = dd[i];
      z.push_back(zz[i]); // only transfer ownership
    } // end for over i
    zz.clear(); // do not free memory -- the memory has been reused in z
  } // end while
  EigenSort(d,z,nPair);
  delete [] dd;
} // end member function PurgeRestart

// devise a vector x, such that it is orthogonal to span{z}, and has increasing
// weight upwards (i.e., the former coordinates have more weight than the latter)
// \param n: the dimension of the matrix to be diagonalized
void TALanczos::InitializeVector(TASparseVec &x, TASparseMatrix &z, int n){
  TARandom r;
  const int nn = n < 50 ? n : 50; // no more than 50 may be a good choice

  x.Clear();
  // Exp(tau): exp(-x/tau)
  for(int i = 0; i < nn; i++) x.Fill(i, r.Gaus((nn-i), 0.5));
  x.normalize(); x.Purify(z);
} // end member function InitializeVector

// note that eigenvalues are not ordered in Jacobi. This routine below sorts
// the input eigenvalues into descending order, and rearranges the eigenvectors
// correspondingly. The method is straight insertion
void TALanczos::EigenSort(double *d, TASparseMatrix &v, int n){
  double p; TASparseVec *vp; // temporary variables for data exchange
  int k; // subscript of the maximum
  for(int i = 0; i < n-1; i++){
    p = d[k=i];
    for(int j = i+1; j < n; j++) if(fabs(d[j]) >= fabs(p)) p = d[k=j];
    if(k != i){
      d[k] = d[i]; d[i] = p; // exchange di and dk
      vp = v[i]; v[i] = v[k]; v[k] = vp; // exchange column i and column k
    } // end if
  } // end for over i
} // end member function EigenSort

#ifdef _ROOT_
#ifdef CATCH_CNOK
TEST_CASE("Lanczos Test - Sparse Matrix", "[.rst]"){
  PARA->SetBaseDir("config/fci");
  TAHPairing *p = new TAHPairing("ham.yaml");
  const int n = int(p->GetNBasis());
  TASparseMatrix v; double d[n];
  for(int i = 0; i < n; i++) d[i] = -9999.;
  // TASPStateManager::Instance()->Print(); // DEBUG
  // TAMBSDManager::Instance()->Print(); // DEBUG

  TALanczos::PurgeRestart(*p, d, v, n);
  CHECK(d[0] == Approx(2.).epsilon(1.e-15));
  CHECK(d[1] == Approx(2.).epsilon(1.e-15));
  CHECK(d[2] == Approx(2.).epsilon(1.e-15));
  CHECK(d[3] == Approx(0.).epsilon(1.e-15));
  CHECK(d[4] == Approx(0.).epsilon(1.e-15));
  CHECK(d[5] == Approx(0.).epsilon(1.e-15));
  CHECK(d[6] == Approx(0.).epsilon(1.e-15));
  CHECK(d[7] == Approx(0.).epsilon(1.e-15));
  CHECK(d[8] == Approx(0.).epsilon(1.e-15));
  // orthogonality check of the eigenvectors //
  CHECK(v[0]->DotProduct(*v[1]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[2]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[3]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[4]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[5]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[6]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[7]) == Approx(0.).margin(1.e-7));
  CHECK(v[0]->DotProduct(*v[8]) == Approx(0.).margin(1.e-7));
  CHECK(v[2]->DotProduct(*v[0]) == Approx(0.).margin(1.e-7));
  CHECK(v[3]->DotProduct(*v[1]) == Approx(0.).margin(1.e-7));
  CHECK(v[4]->DotProduct(*v[2]) == Approx(0.).margin(1.e-7));
  CHECK(v[5]->DotProduct(*v[3]) == Approx(0.).margin(1.e-7));
  CHECK(v[6]->DotProduct(*v[4]) == Approx(0.).margin(1.e-7));
  CHECK(v[7]->DotProduct(*v[5]) == Approx(0.).margin(1.e-7));
  CHECK(v[8]->DotProduct(*v[2]) == Approx(0.).margin(1.e-7));

  v.Save();
	TATreeCol::CloseFile();
} // end TEST_CASE .rst
#endif
#endif
